package com.aij.starter.stream;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.LinkedBlockingQueue;

import org.springframework.util.StringUtils;

import com.aij.starter.event.ChatStreamEventHandler;
import com.aij.starter.stream.StreamClient.StreamCallback;
import com.alibaba.fastjson.JSON;

/**
 * 对接ChatGLM Stream接口
 * @author aij
 */
public class ChatGlmStream {
    
    public static class Message {
    	public String role;
    	public String content;
    }
    
    public static class Request {
    	public String model;
    	public Boolean stream;
    	public List<Message> messages;
    }
    
    public static class Choice {
    	public Message delta;
    }
    
    public static class Response {
    	public List<Choice> choices;
    }
	
    public interface OnNewWord {
        void callback(String text);
    }

    static String url = "";
    static StreamClient client = null;
    static Thread thread = null;
    static LinkedBlockingQueue<Request> queue = new LinkedBlockingQueue<Request>();
    static boolean isRunning = false;
    public static void Initialize(String model)
    {
    	url = model;
    	client = new StreamClient();
        thread = new Thread(new StreamRunnable());
        thread.start();
        isRunning = true;
    }
    public static void Initialize(String model, String prompt)
    {
    	Initialize(model);
    }
    
    private static class StreamRunnable implements Runnable {
    	 
        public void run()
        {
        	while (isRunning) {
                if (queue.size() > 0)
                {
                	Request req = queue.poll();
            		client.postEventStream(url, req, new StreamCallback() {
            			@Override
                        public void callback(String txt) {
            				if (StringUtils.isEmpty(txt) || "[DONE]".equals(txt) || !txt.startsWith("{")) return;
            				try {
	            	    		Response res = JSON.parseObject(txt, Response.class);
	            	    		List<Choice> choices = res.choices;
	            	    		if (choices != null && !choices.isEmpty()) {
	            	    			String message = choices.get(0).delta.content;
	            	    	        //System.out.println(message);
	            		        	ChatStreamEventHandler handler = new ChatStreamEventHandler();
	            		        	handler.callback(message);
	            	    		}
            				} catch (Exception e) {
            					e.printStackTrace();
            				}
            			}
                	});
                }
                try {
					Thread.sleep(1);
				} catch (InterruptedException e) {
					e.printStackTrace();
				}
        	}
        	
        }
    }

    public static void Run(String word)
    {
    	Message msg = new Message();
    	msg.content = word;
    	msg.role = "user";
    	List<Message> msgs = new ArrayList<>();
    	msgs.add(msg);
    	Request req = new Request();
    	req.stream = true;
    	req.model = "chatglm2-6b";
    	req.messages = msgs;
    	queue.add(req);
    }

    public static void Destroy()
    {
        isRunning = false;
    	if (thread != null) {
    		try {
				thread.join();
				thread = null;
		    	client = null;
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
    	}
    }
 
	public static void main(String[] args) {
		String url = "http://192.168.0.100:8000/v1/chat/completions";
		ChatGlmStream.Initialize(url);
		ChatGlmStream.Run("你好");
        try {
			Thread.sleep(5);
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
		ChatGlmStream.Destroy();
	}
}
